from pathlib import Path
from PIL import Image
import json
import numpy as np
import random

import torch
import torch.utils.data
import torchvision

import datasets.transforms as T
from .hico import HICODetection as HICO
from .hico import make_transforms

def build(image_set, args):
    root = Path(args.hoi_path)
    assert root.exists(), f'provided HOI path {root} does not exist'
    PATHS = {
        'train': (root / 'images' / 'train2015', root / 'annotations' / 'trainval_hico.json'),
        'val': (root / 'images' / 'test2015', root / 'annotations' / 'test_hico.json')
    }
    img_folder, anno_file = PATHS[image_set]
    dataset = HICODetection(image_set, img_folder, anno_file, 
                            transforms=make_transforms(image_set, args),
                            num_queries=args.num_queries,
                            dataset_file=args.dataset_file)
    #if image_set == 'val':
    dataset.set_rare_hois(PATHS['train'][1])
    CORRECT_MAT_PATH = root / 'annotations' / 'corre_hico.npy'
    dataset.load_correct_mat(CORRECT_MAT_PATH)
    return dataset



class HICODetection(HICO):

    def __init__(self, img_set, img_folder, anno_file, transforms, num_queries, dataset_file):
        super(HICODetection, self).__init__(img_set, img_folder, anno_file, transforms, num_queries, dataset_file)

    def get_hoi(self, img_anno, target):
        if self.img_set == 'train':
            kept_box_indices = [label[0] for label in target['labels']]
            
            target['labels'] = target['labels'][:, 1]
            obj_labels, verb_labels, sub_boxes, obj_boxes = [], [], [], []
            
            sub_obj_pairs = []
            for hoi in img_anno['hoi_annotation']:  #   一个动作交互
                if hoi['subject_id'] not in kept_box_indices or \
                  (hoi['object_id'] != -1 and hoi['object_id'] not in kept_box_indices):
                    continue
                sub_obj_pair = (hoi['subject_id'], hoi['object_id'])
                if sub_obj_pair in sub_obj_pairs:
                    # verb是多目标
                    verb_labels[sub_obj_pairs.index(sub_obj_pair)][self._valid_verb_ids.index(hoi['category_id'])] = 1
                else:
                    sub_obj_pairs.append(sub_obj_pair)
                    
                    #   obj_labels(kh, )
                    if hoi['object_id'] == -1:
                        obj_labels.append(torch.tensor(len(self._valid_obj_ids)))   #80-即没有物体
                    else:
                        obj_labels.append(target['labels'][kept_box_indices.index(hoi['object_id'])])
                    
                    # #   verb_labels(kh, 29)
                    verb_label = [0 for _ in range(len(self._valid_verb_ids))]
                    verb_label[self._valid_verb_ids.index(hoi['category_id'])] = 1
                    verb_labels.append(verb_label)
                    
                    #   sub_boxes(kh, 4)
                    sub_box = target['boxes'][kept_box_indices.index(hoi['subject_id'])]
                    sub_boxes.append(sub_box)
                    
                    #   obj_boxes
                    if hoi['object_id'] == -1:
                        #obj_box = torch.zeros((4,), dtype=torch.float32)#没有物体bbox给4个零
                        obj_box = sub_box.clone()
                    else:
                        obj_box = target['boxes'][kept_box_indices.index(hoi['object_id'])]
                    obj_boxes.append(obj_box)
                    
################################################################################################################   
            have_hoi = len(sub_boxes)
            #   没有交互的物体对
            for _, sub_obj_pair in enumerate(sub_obj_pairs):
                sub_id  = sub_obj_pair[0]
                #   其他物体
                for _, sub_obj_pair_obj in enumerate(sub_obj_pairs):
                    obj_id  = sub_obj_pair_obj[1]
                    sub_obj_pair = (sub_id, obj_id)
                    if obj_id != -1 and \
                    sub_obj_pair not in sub_obj_pairs:
                        sub_obj_pairs.append(sub_obj_pair)
                        
                        obj_labels.append(target['labels'][kept_box_indices.index(obj_id)])
                        
                        verb_label = [0 for _ in range(len(self._valid_verb_ids))]
                        verb_labels.append(verb_label)
                        
                        sub_box = target['boxes'][kept_box_indices.index(sub_id)]
                        sub_boxes.append(sub_box)
                        
                        obj_box = target['boxes'][kept_box_indices.index(obj_id)]
                        obj_boxes.append(obj_box)

                #   自己
                if self.dataset_file=='vcoco':
                    sub_obj_pair = (sub_id, -1)
                    if sub_obj_pair not in sub_obj_pairs:
                        sub_obj_pairs.append(sub_obj_pair)
                        
                        obj_labels.append(torch.tensor(len(self._valid_obj_ids))) 
                        
                        verb_label = [0 for _ in range(len(self._valid_verb_ids))]
                        verb_labels.append(verb_label)
                        
                        sub_box = target['boxes'][kept_box_indices.index(sub_id)]
                        sub_boxes.append(sub_box)

                        #obj_box = torch.zeros((4,), dtype=torch.float32)#没有物体bbox给4个零
                        obj_box = target['boxes'][kept_box_indices.index(sub_id)]
                        obj_boxes.append(obj_box)   
             
            if len(sub_obj_pairs) != 0:
                obj_labels =  torch.stack(obj_labels)
                verb_labels = torch.as_tensor(verb_labels, dtype=torch.float32)
                sub_boxes = torch.stack(sub_boxes)
                obj_boxes =  torch.stack(obj_boxes)
                
                if  len(sub_obj_pairs) >= self.num_queries:
                    x = list(range(have_hoi))
                    y = list(range(have_hoi, len(sub_obj_pairs)))
                    random.shuffle(y)
                    index = x + y
                    index = index[:self.num_queries]
                    
                    obj_labels = obj_labels[index]
                    verb_labels = verb_labels[index]
                    sub_boxes = sub_boxes[index]
                    obj_boxes = obj_boxes[index]
                matching_labels = torch.ones_like(obj_labels, dtype=torch.int64)
                #matching_labels[have_hoi:] = 0.
                
                
            if len(sub_obj_pairs) == 0:
                target['obj_labels'] = torch.zeros((0,), dtype=torch.int64)
                target['verb_labels'] = torch.zeros((0, len(self._valid_verb_ids)), dtype=torch.float32)
                target['sub_boxes'] = torch.zeros((0, 4), dtype=torch.float32)
                target['obj_boxes'] = torch.zeros((0, 4), dtype=torch.float32)
                target['matching_labels'] = torch.zeros((0,),  dtype=torch.int64)
            else:
                target['obj_labels']  = obj_labels
                target['verb_labels'] = verb_labels
                target['sub_boxes']   = sub_boxes
                target['obj_boxes']   = obj_boxes
                target['matching_labels'] = matching_labels
        else:
            hois = []
            for hoi in img_anno['hoi_annotation']:
                hois.append((hoi['subject_id'], 
                             hoi['object_id'], 
                             self._valid_verb_ids.index(hoi['category_id'])))
            target['hois'] = torch.as_tensor(hois, dtype=torch.int64) 
                   
        return target